package com.example.testwebsockert;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import jakarta.websocket.*;
import jakarta.websocket.server.PathParam;
import jakarta.websocket.server.ServerEndpoint;
import org.springframework.stereotype.Component;
import java.io.IOException;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

@Component
@ServerEndpoint(value = "/websocket/{clientId}")
public class SocketController {

    // 存储所有连接的客户端
    private static final Map<String, Session> clients = new ConcurrentHashMap<>();
    private static final Map<String, String> userId = new ConcurrentHashMap<>();
    private static final Map<String, List<String>> offMessage = new ConcurrentHashMap<>();
    private static final Set<String> users = ConcurrentHashMap.newKeySet();

    // 分片消息存储
    private static final Map<String, StringBuilder> messageChunks = new ConcurrentHashMap<>();
    private static final Map<String, Integer> expectedChunks = new ConcurrentHashMap<>();
    private static final Map<String, Integer> receivedChunks = new ConcurrentHashMap<>();

    // 异步任务执行器
    private static final ExecutorService executor = Executors.newSingleThreadExecutor();

    // 消息大小限制常量 (5MB)
    private static final long MAX_TEXT_MESSAGE_SIZE = 5 * 1024 * 1024;

    @OnOpen
    public void onOpen(Session session, @PathParam("clientId") String clientId) {
        clients.put(clientId, session);

        // 检查离线消息
        if (users.contains(clientId)) {
            List<String> messages = offMessage.getOrDefault(clientId, Collections.emptyList());
            if (!messages.isEmpty()) {
                executor.submit(() -> {
                    try {
                        for (String msg : messages) {
                            sendMessageToClient(session, msg);
                            Thread.sleep(50);
                        }
                    } catch (InterruptedException e) {
                        Thread.currentThread().interrupt();
                    }
                });
                users.remove(clientId);
                offMessage.remove(clientId);
            }
        }
    }

    @OnClose
    public void onClose(Session session, @PathParam("clientId") String clientId) {
        clients.remove(clientId);
        userId.remove(clientId);
        users.add(clientId);
        cleanupChunkedMessages(clientId);
    }

    // 关键修改：通过注解设置消息大小限制
    @OnMessage(maxMessageSize = 5 * 1024 * 1024) // 5MB
    public void onMessage(String message, Session session, @PathParam("clientId") String clientId) {
        try {
            // 检查消息大小
            if (message.length() > MAX_TEXT_MESSAGE_SIZE) {
                handleError(session, "消息过大，最大支持5MB");
                return;
            }

            JSONObject jsonObject = JSON.parseObject(message);
            String type = jsonObject.getString("type");

            // 处理分片消息
            if ("largeMessageStart".equals(type)) {
                String messageId = jsonObject.getString("messageId");
                int totalChunks = jsonObject.getIntValue("totalChunks");
                messageChunks.put(messageId, new StringBuilder());
                expectedChunks.put(messageId, totalChunks);
                receivedChunks.put(messageId, 0);
                return;
            }
            else if ("largeMessageChunk".equals(type)) {
                String messageId = jsonObject.getString("messageId");
                int chunkIndex = jsonObject.getIntValue("chunkIndex");
                String chunkData = jsonObject.getString("data");

                if (messageChunks.containsKey(messageId)) {
                    messageChunks.get(messageId).append(chunkData);
                    receivedChunks.put(messageId, receivedChunks.get(messageId) + 1);
                }
                return;
            }
            else if ("largeMessageEnd".equals(type)) {
                String messageId = jsonObject.getString("messageId");

                if (messageChunks.containsKey(messageId) &&
                        receivedChunks.getOrDefault(messageId, 0) == expectedChunks.getOrDefault(messageId, 0)) {

                    String fullMessage = messageChunks.get(messageId).toString();
                    processCompleteMessage(fullMessage, session, clientId);
                    cleanupChunkedMessage(messageId);
                }
                return;
            }

            // 处理普通消息
            processCompleteMessage(message, session, clientId);

        } catch (Exception e) {
            handleError(session, "处理消息错误: " + e.getMessage());
        }
    }

    private void processCompleteMessage(String message, Session session, String clientId) {

        JSONObject jsonObject = JSON.parseObject(message);
        String type = jsonObject.getString("type");
        if ("message".equals(type)) {
            // 存储离线消息
            executor.submit(() -> {
                for (String user : users) {
                    offMessage.computeIfAbsent(user, k -> new ArrayList<>()).add(message);
                }
            });
            userId.put(clientId, jsonObject.getString("username"));
            sendMessageToAll(clientId, session, message, "message");
        } else if ("leave".equals(type)) {
            userId.remove(clientId);
        }else if("privateMessage".equals(type)){
            JSONObject jsonObject1 = jsonObject.getJSONObject("toUser");
            String toUserId = jsonObject1.get("userid").toString();
            Session session1 = clients.get(toUserId);
            sendMessageToClient(session1,message);
        } else if ("userList".equals(type)) {
            userId.put(clientId, jsonObject.getString("username"));
            List<Map<String,String>> userList = new ArrayList<>();
            Map<String, Object> map = new ConcurrentHashMap<>();
            map.put("type", "userList");
            userId.forEach((k, v) -> {
                Map<String,String> user = new ConcurrentHashMap<>();
                user.put("username", v);
                user.put("userid", k);
                userList.add(user);
            });
            map.put("data", userList);
            sendMessageToAll(clientId, session, JSON.toJSONString(map), "userList");
        } else {
            sendMessageToAll(clientId, session, message, "userList");
        }
    }

    private void sendMessageToClient(Session session, String message) {
        try {
            if (message.getBytes().length > MAX_TEXT_MESSAGE_SIZE) {
                handleError(session, "消息过大，最大支持5MB");
                return;
            }

            if (session.isOpen()) {
                session.getBasicRemote().sendText(message);
            }
        } catch (IOException e) {
            handleError(session, "发送消息失败: " + e.getMessage());
        }
    }

    public static void sendMessageToAll(String clientId, Session my, String message, String type) {
        if (message.getBytes().length > MAX_TEXT_MESSAGE_SIZE) {
            try {
                ((SocketController) my.getUserProperties().get("handler")).handleError(my, "消息过大，最大支持5MB");
            } catch (Exception e) {
                e.printStackTrace();
            }
            return;
        }

        if ("userList".equals(type)) {
            clients.forEach((k, v) -> {
                try {
                    if (v.isOpen()) {
                        v.getBasicRemote().sendText(message);
                    }
                } catch (IOException e) {
                    e.printStackTrace();
                }
            });
        } else {
            clients.forEach((k, v) -> {
                try {
                    if (!k.equals(clientId) && v.isOpen()) {
                        v.getBasicRemote().sendText(message);
                    }
                } catch (IOException e) {
                    e.printStackTrace();
                }
            });
        }
    }

    private void handleError(Session session, String errorMessage) {
        try {
            if (session.isOpen()) {
                Map<String, Object> errorResponse = new HashMap<>();
                errorResponse.put("type", "error");
                errorResponse.put("message", errorMessage);
                session.getBasicRemote().sendText(JSON.toJSONString(errorResponse));
            }
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    private void cleanupChunkedMessage(String messageId) {
        messageChunks.remove(messageId);
        expectedChunks.remove(messageId);
        receivedChunks.remove(messageId);
    }

    private void cleanupChunkedMessages(String clientId) {
        // 清理相关分片消息
    }

    @OnError
    public void onError(Session session, Throwable throwable, @PathParam("clientId") String clientId) {
        String errorMsg = throwable.getMessage();
        if (errorMsg != null && (errorMsg.contains("too big") || errorMsg.contains("过大"))) {
            handleError(session, "消息过大，最大支持5MB");
        } else {
            handleError(session, "WebSocket错误: " + errorMsg);
        }
    }
}
